package com.arvin.faceDetect.utils

import android.graphics.Bitmap
import android.graphics.BitmapFactory
import android.graphics.ImageFormat
import android.graphics.Matrix
import android.graphics.Rect
import android.graphics.YuvImage
import androidx.camera.core.ImageProxy
import java.io.ByteArrayOutputStream
import java.nio.ByteBuffer

fun ImageProxy.toBitmap(): Bitmap {
    val image = this.image ?: return Bitmap.createBitmap(1, 1, Bitmap.Config.ARGB_8888)
    
    when (image.format) {
        ImageFormat.YUV_420_888 -> {
            val planes = image.planes
            
            val yPlane = planes[0]
            val uPlane = planes[1]
            val vPlane = planes[2]

            val yBuffer = yPlane.buffer
            val uBuffer = uPlane.buffer
            val vBuffer = vPlane.buffer

            val ySize = yBuffer.remaining()
            val uSize = uBuffer.remaining()
            val vSize = vBuffer.remaining()

            val nv21 = ByteArray(ySize + uSize + vSize)

            // 复制 Y 平面数据
            yBuffer.get(nv21, 0, ySize)
            
            // 交错复制 U 和 V 平面数据
            val uvPixelStride = uPlane.pixelStride
            if (uvPixelStride == 1) {
                uBuffer.get(nv21, ySize, uSize)
                vBuffer.get(nv21, ySize + uSize, vSize)
            } else {
                // 交错的 UV 平面，需要逐像素复制
                var pos = ySize
                while (uBuffer.hasRemaining()) {
                    nv21[pos++] = vBuffer.get()
                    nv21[pos++] = uBuffer.get()
                }
            }

            val yuvImage = YuvImage(nv21, ImageFormat.NV21, width, height, null)
            val out = ByteArrayOutputStream()
            yuvImage.compressToJpeg(Rect(0, 0, width, height), 100, out)
            val imageBytes = out.toByteArray()
            var bitmap = BitmapFactory.decodeByteArray(imageBytes, 0, imageBytes.size)

            // 处理图像旋转
            if (imageInfo.rotationDegrees != 0) {
                val matrix = Matrix()
                matrix.postRotate(imageInfo.rotationDegrees.toFloat())
                bitmap = Bitmap.createBitmap(
                    bitmap, 0, 0, bitmap.width, bitmap.height, matrix, true
                )
            }

            return bitmap
        }
        else -> {
            // 对于其他格式，返回一个空白位图
            return Bitmap.createBitmap(1, 1, Bitmap.Config.ARGB_8888)
        }
    }
} 